#!/bin/bash

# Ubuntu桌面环境 - 存储空间清理脚本

set -e

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m'

print_message() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

print_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

print_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

print_header() {
    echo -e "${BLUE}================================${NC}"
    echo -e "${BLUE}  存储空间清理脚本${NC}"
    echo -e "${BLUE}================================${NC}"
}

# 检查Docker
check_docker() {
    if ! command -v docker &> /dev/null; then
        print_error "Docker未安装"
        exit 1
    fi
    
    if ! docker info &> /dev/null; then
        print_error "Docker服务未运行"
        exit 1
    fi
}

# 显示清理前的状态
show_before_status() {
    print_message "清理前的存储状态："
    echo
    
    echo "系统磁盘使用情况："
    df -h
    
    echo
    echo "Docker存储使用情况："
    docker system df
    
    echo
    echo "Docker镜像："
    docker images --format "table {{.Repository}}\t{{.Tag}}\t{{.Size}}\t{{.CreatedAt}}" | head -10
    
    echo
    echo "Docker容器："
    docker ps -a --format "table {{.Names}}\t{{.Status}}\t{{.Size}}" | head -10
    
    echo
    echo "Docker卷："
    docker volume ls --format "table {{.Name}}\t{{.Driver}}" | head -10
}

# 停止所有相关容器
stop_containers() {
    print_message "停止所有相关容器..."
    
    # 停止Ubuntu桌面环境容器
    docker-compose -f docker-compose.lightweight.yml down 2>/dev/null || true
    docker-compose -f docker-compose.yml down 2>/dev/null || true
    docker-compose -f docker-compose-novnc.yml down 2>/dev/null || true
    
    # 停止其他可能相关的容器
    docker stop $(docker ps -q --filter "name=ubuntu-desktop") 2>/dev/null || true
    docker stop $(docker ps -q --filter "name=novnc") 2>/dev/null || true
    
    print_message "容器已停止"
}

# 清理Docker系统
cleanup_docker_system() {
    print_message "清理Docker系统..."
    
    # 清理未使用的镜像
    print_message "清理未使用的镜像..."
    docker image prune -a -f
    
    # 清理未使用的容器
    print_message "清理未使用的容器..."
    docker container prune -f
    
    # 清理未使用的卷
    print_message "清理未使用的卷..."
    docker volume prune -f
    
    # 清理未使用的网络
    print_message "清理未使用的网络..."
    docker network prune -f
    
    # 清理构建缓存
    print_message "清理构建缓存..."
    docker builder prune -a -f
    
    # 清理整个系统
    print_message "深度清理Docker系统..."
    docker system prune -a -f --volumes
    
    print_message "Docker系统清理完成"
}

# 清理特定镜像
cleanup_specific_images() {
    print_message "清理特定镜像..."
    
    # 删除Ubuntu桌面环境相关镜像
    docker rmi ubuntu-desktop:lightweight-xfce-zh 2>/dev/null || true
    docker rmi ubuntu-desktop:xfce-zh 2>/dev/null || true
    docker rmi ubuntu-desktop:lightweight-vnc 2>/dev/null || true
    
    # 删除悬空镜像
    docker rmi $(docker images -f "dangling=true" -q) 2>/dev/null || true
    
    print_message "特定镜像清理完成"
}

# 清理系统临时文件
cleanup_system_temp() {
    print_message "清理系统临时文件..."
    
    # 清理apt缓存
    if command -v apt-get &> /dev/null; then
        print_message "清理APT缓存..."
        sudo apt-get clean
        sudo apt-get autoremove -y
    fi
    
    # 清理临时文件
    print_message "清理临时文件..."
    sudo rm -rf /tmp/* 2>/dev/null || true
    sudo rm -rf /var/tmp/* 2>/dev/null || true
    
    # 清理日志文件
    print_message "清理日志文件..."
    sudo find /var/log -name "*.log" -type f -size +100M -delete 2>/dev/null || true
    sudo find /var/log -name "*.gz" -type f -delete 2>/dev/null || true
    
    print_message "系统临时文件清理完成"
}

# 清理项目目录
cleanup_project_dir() {
    print_message "清理项目目录..."
    
    # 清理用户数据和共享文件夹
    if [ -d "user-data" ]; then
        print_message "清理用户数据目录..."
        rm -rf user-data/*
        mkdir -p user-data
    fi
    
    if [ -d "shared" ]; then
        print_message "清理共享文件夹..."
        rm -rf shared/*
        mkdir -p shared
    fi
    
    # 清理Docker相关文件
    print_message "清理Docker相关文件..."
    rm -f .env 2>/dev/null || true
    
    print_message "项目目录清理完成"
}

# 显示清理后的状态
show_after_status() {
    print_message "清理后的存储状态："
    echo
    
    echo "系统磁盘使用情况："
    df -h
    
    echo
    echo "Docker存储使用情况："
    docker system df
    
    echo
    echo "释放的空间："
    echo "请对比清理前后的数据"
}

# 显示清理建议
show_cleanup_suggestions() {
    echo
    echo -e "${GREEN}================================${NC}"
    echo -e "${GREEN}  清理完成！${NC}"
    echo -e "${GREEN}================================${NC}"
    echo
    
    echo "已执行的清理操作："
    echo "✅ 停止所有相关容器"
    echo "✅ 清理Docker镜像、容器、卷、网络"
    echo "✅ 清理构建缓存"
    echo "✅ 清理系统临时文件"
    echo "✅ 清理项目目录"
    echo
    
    echo "后续建议："
    echo "1. 定期运行此脚本清理空间"
    echo "2. 使用轻量级版本减少空间占用"
    echo "3. 定期清理不需要的Docker镜像"
    echo "4. 监控磁盘使用情况"
    echo
    
    echo "重新启动服务："
    echo "  - 轻量级版本: docker-compose -f docker-compose.lightweight.yml up -d"
    echo "  - noVNC版本: ./start-novnc.sh"
    echo "  - 标准版本: docker-compose up -d"
    echo
}

# 主函数
main() {
    print_header
    
    # 检查Docker
    check_docker
    
    # 显示清理前的状态
    show_before_status
    
    echo
    read -p "是否继续清理？这将停止所有相关容器并清理大量数据 (y/N): " -n 1 -r
    echo
    if [[ ! $REPLY =~ ^[Yy]$ ]]; then
        print_message "用户取消操作"
        exit 0
    fi
    
    # 执行清理操作
    stop_containers
    cleanup_docker_system
    cleanup_specific_images
    cleanup_system_temp
    cleanup_project_dir
    
    # 显示清理后的状态
    show_after_status
    
    # 显示清理建议
    show_cleanup_suggestions
}

# 捕获信号
trap 'print_error "脚本被中断"; exit 1' INT TERM

# 运行主函数
main "$@" 